"""
pip install kokoro-onnx soundfile

wget https://github.com/thewh1teagle/kokoro-onnx/releases/download/model-files-v1.0/kokoro-v1.0.onnx
wget https://github.com/thewh1teagle/kokoro-onnx/releases/download/model-files-v1.0/voices-v1.0.bin
python examples/with_session.py
"""

import soundfile as sf
from kokoro_onnx import Kokoro
from onnxruntime import InferenceSession
import onnxruntime
import os


def create_session():
    # See list of providers https://github.com/microsoft/onnxruntime/issues/22101#issuecomment-2357667377
    providers = onnxruntime.get_available_providers()
    print(f"Available onnx runtime providers: {providers}")

    # See session options https://onnxruntime.ai/docs/performance/tune-performance/threading.html#thread-management
    sess_options = onnxruntime.SessionOptions()
    # Set threads to num of CPU cores
    cpu_count = os.cpu_count()
    print(f"Setting threads to CPU cores count: {cpu_count}")
    sess_options.intra_op_num_threads = cpu_count
    session = InferenceSession(
        "kokoro-v1.0.onnx", providers=providers, sess_options=sess_options
    )
    return session


session = create_session()
kokoro = Kokoro.from_session(session, "voices-v1.0.bin")
samples, sample_rate = kokoro.create(
    "Hello. This audio generated by kokoro!", voice="af_sarah", speed=1.0, lang="en-us"
)
sf.write("audio.wav", samples, sample_rate)
print("Created audio.wav")
